#include "../DWMRICore/DWI.h"
#include "../DWMRICore/DTI.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/TwoTensor.h"

#include <levmar/levmar.h>

#include <teem/ell.h>
#include <teem/ten.h>

#define _USE_MATH_DEFINES
#include <math.h>
#include <time.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <omp.h>


char *pDWIPathname = NULL;
char *pDTIPathname = NULL;
char *pModelSelectionPathname = NULL;
char *twoTensorPathname = NULL;

int scale = 1;

/* MLS fiber tracking parameters */
float MLSSize = 2.0f;
CPoint3F *pSamplePoints = NULL;
float *pSamplePointsWeight = NULL;

CDWI *pDWI = NULL;
CDTI *pDTI = NULL;
CScalar *pModelSelection = NULL;

struct AData {
	float s0;
	float bValue;
	CVector3F *gradientOrientations;
	float lam3;
};


void PrintUsage()
{
	printf("Usage: TwoTensorFitting -dwi dwi_file -dti dti_file -modelselection model_selection_file -dst dst_file [-thresh threshold]\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dwi") == 0) {
			pDWIPathname = argv[++i];
		} else if (strcmp(argv[i], "-dti") == 0) {
			pDTIPathname = argv[++i];
		} else if (strcmp(argv[i], "-modelselection") == 0) {
			pModelSelectionPathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			twoTensorPathname = argv[++i];
		} else {
			return -1;
		}
	}

	if (pDWIPathname == NULL || pDTIPathname == NULL || 
		pModelSelectionPathname == NULL || twoTensorPathname == NULL)
		return -2;

	return 0;
}

/* Form a 2D tensor from the parameters */
void PeledRotate2D(double ten[7], double lam1, double lam3, double phi) {
  double cc, ss, d3, d1, d2;

  cc = cos(phi);
  ss = sin(phi);
  d1 = cc*cc*lam1 + ss*ss*lam3;
  d3 = cc*ss*(lam1 - lam3);
  d2 = ss*ss*lam1 + cc*cc*lam3;

  TEN_T_SET(ten, 1.0, d1, d3, 0.0, d2, 0.0, lam3);
  return;
}

/* call back function for levmar optimaization */
void TwoTensorFittingCB(double *pp, double *xx, int mm, int nn, void *_pvlData) {
  /* static const char me[]="_tenLevmarPeledCB"; */
  double tenA[7], tenB[7];
  int ii;
  AData *pvlData;
  CVector3F *gradients;
 
  AIR_UNUSED(mm);
  pvlData = AIR_CAST(AData *, _pvlData);

  /* Form the tensors using the estimated parms */
  PeledRotate2D(tenA, pp[0], pvlData->lam3, pp[2]);
  PeledRotate2D(tenB, pp[0], pvlData->lam3, pp[3]);

  gradients = AIR_CAST(CVector3F *, pvlData->gradientOrientations);

  for (ii = 0; ii < nn; ii++) {
    double argA, argB, sigA, sigB;
	double egrad[3];
	egrad[0] = gradients[ii].m_x;
	egrad[1] = gradients[ii].m_y;
	egrad[2] = gradients[ii].m_z;
	argA = -pvlData->bValue * TEN_T3V_CONTR(tenA, egrad);
	argB = -pvlData->bValue * TEN_T3V_CONTR(tenB, egrad);
	sigA = exp(argA);
	sigB = exp(argB);

    xx[ii] = pp[1] * sigA + (1.0 - pp[1]) * sigB;
  }
  return;
}

void Fitting(float *_dwi, const float s0, const float bValue, 
	float *eval, float *evec, CVector3F *gradients, const int n, float *twoTensor)
{
	double guess[4], loBnd[4], upBnd[4], opts[LM_OPTS_SZ], info[LM_INFO_SZ];
	double tenA[7], tenB[7];
	double matA[9], matB[9], rott[9], matTemp[9];
	double *dwi = new double[n];

	// dwi
	for (int i = 0; i < n; ++i) {
		dwi[i] = _dwi[i] / s0;
	}

	// set lower and upper bound
	loBnd[0] = 0.0;
	loBnd[1] = 0.0;
	loBnd[2] = -M_PI / 2.0;
	loBnd[3] = -M_PI / 2.0;
	upBnd[0] = eval[0] * 5.0;
	upBnd[1] = 1.0;
	upBnd[2] = M_PI / 2.0;
	upBnd[3] = M_PI / 2.0;

	// set initial guess
	guess[0] = eval[0];
	guess[1] = 0.5;
	guess[2] = M_PI / 4.0;
	guess[3] = -M_PI / 4.0;

	// set constraints for LM optimization
	opts[0] = 1E-3;
	opts[1] = 1E-8;
	opts[2] = 1E-8;
	opts[3] = 1E-8;
	opts[4] = 1E-8;

	AData aData;
	aData.s0 = s0;
	aData.bValue = bValue;
	aData.gradientOrientations = gradients;
	aData.lam3 = eval[2];

	int work_size = LM_BC_DIF_WORKSZ(4, n);
	double *work = new double[work_size*2];
	int lmret = dlevmar_bc_dif(TwoTensorFittingCB, guess, dwi, 4, n, 
		loBnd, upBnd, NULL, 1000, opts, info, NULL, NULL, (void *)&aData);
	SafeDeleteArray(work);

	SafeDeleteArray(dwi);

	if (-1 == lmret) {
		printf("dlevmar_bc_dif() failed!\n");
	} else {
		// get two tensor fit result
		PeledRotate2D(tenA, guess[0], eval[2], guess[2]);
		PeledRotate2D(tenB, guess[0], eval[2], guess[3]);
		TEN_T2M(matA, tenA);
		TEN_T2M(matB, tenB);

		ELL_3M_TRANSPOSE(rott, evec);
		ELL_3M_MUL(matTemp, matA, evec);
		ELL_3M_MUL(matA, rott, matTemp);
		ELL_3M_MUL(matTemp, matB, evec);
		ELL_3M_MUL(matB, rott, matTemp);

		TEN_M2T(tenA, matA);
		TEN_M2T(tenB, matB);

		if (guess[1] > 0.5) {
			twoTensor[0] = 1.0f;
			twoTensor[7] = 1.0f - guess[1];

			for (int i = 1; i < 7; ++i) {
				twoTensor[i] = tenA[i];
				twoTensor[i+7] = tenB[i];
			}

		} else {
			twoTensor[0] = 1.0f;
			twoTensor[7] = guess[1];

			for (int i = 1; i < 7; ++i) {
				twoTensor[i] = tenB[i];
				twoTensor[i+7] = tenA[i];
			}
		}
		
	}
}

void DTIMLSFilter(const CPoint3F pos, float *lastDTI, float *filteredDTI)
{
	float m[9];
	TEN_T2M(m, lastDTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	memset(filteredDTI, 0, sizeof(float)*7);
	float temp_DTI[7];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pModelSelection->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		pDTI->GetDTI(samples[i], temp_DTI);

		for (uint j = 0; j < 7; ++j) {
			filteredDTI[j] += pSamplePointsWeight[i] * temp_DTI[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < 7; ++j) {
		filteredDTI[j] = filteredDTI[j] / weight;
	}
}

void DWIMLSFilter(const CPoint3F pos, float *filteredDTI, float *filteredDWI, float *filteredS0)
{
	float m[9];
	TEN_T2M(m, filteredDTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	uint n = pDWI->GetN() - pDWI->GetM();
	memset(filteredDWI, 0, sizeof(float)*n);
	*filteredS0 = 0.0f;
	float *temp_DWI = new float[n];
	float temp_S0 = 0.0f;
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pModelSelection->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		pDWI->GetNS0(samples[i], temp_DWI);
		temp_S0 = pDWI->GetS0(samples[i]);

		for (uint j = 0; j < n; ++j) {
			filteredDWI[j] += pSamplePointsWeight[i] * temp_DWI[j];
		}
		*filteredS0 += pSamplePointsWeight[i] * temp_S0;
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < n; ++j) {
		filteredDWI[j] = filteredDWI[j] / weight;
	}
	*filteredS0 = *filteredS0 / weight;

	SafeDeleteArray(temp_DWI);
}

void TwoTensorFitting()
{
	uint w = pDWI->GetWidth();
	uint h = pDWI->GetHeight();
	uint d = pDWI->GetDepth();
	uint n = pDWI->GetN();
	uint m = pDWI->GetM();
	float bValue = pDWI->GetBValue();

	int sw = w * scale;
	int sh = h * scale;
	int sd = d * scale;

	CVector3F *diffusionGradients = pDWI->GetDiffusionGradients();

	float *twoTensor = new float[sw*sh*sd*14];
	memset(twoTensor, 0, sizeof(float)*sw*sh*sd*14);
		
	for (int z = 0; z < sd; ++z) {
//		omp_set_num_threads(16);
//#pragma omp parallel for
		for (int y = 0; y < sh; ++y) {
			for (int x = 0; x < sw; ++x) {
				int i = (z * sh + y) * sw + x;

				CPoint3F pos;
				pos.m_x = (float)x/(float)scale;
				pos.m_y = (float)y/(float)scale;
				pos.m_z = (float)z/(float)scale;

				float initDTI[7] = {1, 1, 0, 0, 1, 0, 1};
				float filteredDTI[7];
				DTIMLSFilter(pos, initDTI, filteredDTI);

				if (pModelSelection->GetScalar(pos) > 0.5f) {
					float *filteredDWI = new float[n-m];
					float filteredS0 = 0.0f;
					DWIMLSFilter(pos, filteredDTI, filteredDWI, &filteredS0);

					float eval[3], evec[9];
					pDTI->DTI2Eigen(filteredDTI, eval, evec);

					// rotate the graidents
					CVector3F *gradients = new CVector3F[n-m];
					for (uint k = m; k < n; ++k) {
						float v[3], vv[3];
						v[0] = diffusionGradients[k].m_x;
						v[1] = diffusionGradients[k].m_y;
						v[2] = diffusionGradients[k].m_z;
						ELL_3MV_MUL(vv, evec, v);
						gradients[k-m].m_x = vv[0];
						gradients[k-m].m_y = vv[1];
						gradients[k-m].m_z = vv[2];
					}

					// do the do tensor fitting
					float res[14];
					Fitting(filteredDWI, filteredS0, bValue, eval, evec, gradients, n-m, res);
					memcpy(&(twoTensor[i*14]), res, sizeof(float)*14);

					SafeDeleteArray(gradients);
					SafeDeleteArray(filteredDWI);
				} else {
					memcpy(&(twoTensor[i*14]), filteredDTI, sizeof(float)*7);
					twoTensor[i*14] = 1.0f;
				}
			}
		}
		float ff = float(z + 1) / (float)sd * 100.0f;
		printf("\r%.2f%% done.", ff);
	}
	printf("\r100.0%% done.\n");

	FILE *fp = fopen(twoTensorPathname, "wb");
	fwrite(twoTensor, sizeof(float), sw*sh*sd*14, fp);
	fclose(fp);


	SafeDeleteArray(twoTensor);
}

void CreateMLSSamplePoints()
{
	SafeDeleteArray(pSamplePoints);
	SafeDeleteArray(pSamplePointsWeight);
	pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	pSamplePoints[0].m_x = 0.0f;
	pSamplePoints[0].m_y = 0.0f;
	pSamplePoints[0].m_z = 0.0f;
	pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(pSamplePoints[j], pt);
				if (dis < 0.25f) {
					flag = false;
					break;
				}
			}
		}

		pSamplePoints[i] = pt;
		pSamplePointsWeight[i] = InnerProduct(pSamplePoints[i], pSamplePoints[i]) * 4.0f;
		pSamplePointsWeight[i] = expf(pSamplePointsWeight[i] / -2.5f);
	}
}

//int main(int argc, char *argv[])
//{
//	if (ParseArguments(argc, argv) != 0) {
//		PrintUsage();
//		return 0;
//	}
//
//	pDWI = new CDWI();
//	pDWI->OpenDWIFile(pDWIPathname);
//	pDTI = new CDTI();
//	pDTI->OpenDTIFile(pDTIPathname);
//	pModelSelection = new CScalar();
//	pModelSelection->OpenScalarFile(pModelSelectionPathname);
//
//	CreateMLSSamplePoints();
//	TwoTensorFitting();
//
//	SafeDelete(pDWI);
//	SafeDelete(pDTI);
//	SafeDelete(pModelSelection);
//
//	return 0;
//}


int main(int argc, char *argv[])
{
	char *srcPathname = "C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\1X\\TwoTensor_B1000.nhdr";
	char *dstPathname = "C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\1X\\TwoTensorDirection.raw";
	char *faPathname  = "C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\1X\\TwoTensorFA.nrrd";

	CTwoTensor *pTwoTensor = new CTwoTensor();
	pTwoTensor->OpenTwoTensorFile(srcPathname);

	uint w = pTwoTensor->GetWidth();
	uint h = pTwoTensor->GetHeight();
	uint d = pTwoTensor->GetDepth();
	uint n = pTwoTensor->GetN();

	float *directions = new float[w*h*d*6];
	memset(directions, 0, sizeof(float)*w*h*d*6);

	float *fa = new float[w*h*d];
	memset(fa, 0, sizeof(float)*w*h*d);

	for (uint i = 0; i < w*h*d; ++i) {
		float tenA[7], tenB[7];
		pTwoTensor->GetTensorA(i, tenA);
		pTwoTensor->GetTensorB(i, tenB);

		if (tenA[0] > 1e-6) {
			CVector3F dir;
			pTwoTensor->Tensor2Direction(tenA, &dir);
			directions[i*6+0] = dir.m_x;
			directions[i*6+1] = dir.m_y;
			directions[i*6+2] = dir.m_z;
			fa[i] = pTwoTensor->Tensor2FA(tenA);

			if (tenB[0] > 1e-6) {
				CVector3F dir;
				pTwoTensor->Tensor2Direction(tenB, &dir);
				directions[i*6+3] = dir.m_x;
				directions[i*6+4] = dir.m_y;
				directions[i*6+5] = dir.m_z;
				fa[i] += pTwoTensor->Tensor2FA(tenB);
				fa[i] = fa[i] / 2.0f;
			}
		}
	}

	//CScalar *pFA = new CScalar();
	//pFA->CreateScalar(w, h, d, fa);
	//pFA->SaveScalarFile(faPathname);

	FILE *fp = fopen(dstPathname, "wb");
	fwrite(directions, sizeof(float), w*h*d*6, fp);
	fclose(fp);

	SafeDelete(pTwoTensor);
	//SafeDelete(pFA);
	SafeDeleteArray(directions);
	SafeDeleteArray(fa);
}